use crate::common::*;
use futures_util::stream::StreamExt;
use soketto::connection::Sender;
use soketto::handshake::{Client as WsClient, ServerResponse};
use std::net::SocketAddr;
use tokio::net::TcpStream;
use tokio_util::compat::{Compat, TokioAsyncReadCompatExt};
use unmp_link::prelude::*;

pub struct Client {
    identifier: String,
    tx: Sender<Compat<TcpStream>>,
    rx: Rx,
}
impl Client {
    pub async fn new(addr: SocketAddr) -> Result<Connection, ()> {
        let identifier = addr.to_string();
        let socket = TcpStream::connect(addr).await.map_err(|_| ())?;
        let mut client = WsClient::new(socket.compat(), &identifier, PATH);
        client.add_protocol(PROTOCOL);
        let (tx, rx) = match client.handshake().await.map_err(|_| ())? {
            ServerResponse::Accepted { .. } => client.into_builder().finish(),
            ServerResponse::Redirect { .. } => return Err(()),
            ServerResponse::Rejected { .. } => return Err(()),
        };
        Ok(Connection::new(Box::new(Self {
            identifier,
            tx,
            rx: convert_rx(rx),
        })))
    }
}
#[async_trait]
impl ConnectionDriver for Client {
    fn is_subnet(&self) -> bool {
        false
    }
    fn kind(&self) -> &str {
        "ws_client"
    }
    fn id(&self) -> &str {
        &self.identifier
    }
    async fn send(&mut self, buf: Bytes) -> Result<(), SendError> {
        if let Err(_) = self.tx.send_binary(&buf).await {
            return Err(SendError::Disconnected);
        }
        Ok(())
    }
    async fn recv(&mut self) -> Result<Bytes, RecvError> {
        loop {
            let (data_type, data) = match self.rx.next().await {
                Some(Ok(data_type)) => data_type,
                _ => return Err(RecvError::Disconnected),
            };
            if data_type.is_binary() {
                return Ok(data.into());
            }
        }
    }
}
